from core.cuda import cuda
from core.config import no_grad
from core.data.dataloader import DataLoader
from implement.datasource.minist import MNIST
from implement.models.mlp import MLP
from implement.optimizers.adam import Adam
from implement.optimizers.hooks.weight_decay import WeightDecay
from utils.common import accuracy
from utils.functions_collect import softmax_cross_entropy, relu

if '__file__' in globals():
    import os, sys

    sys.path.append(os.path.join(os.path.dirname(__file__), '..'))

max_epoch = 5
batch_size = 100
hidden_size = 1000

train_set = MNIST(train=True)
test_set = MNIST(train=False)
train_loader = DataLoader(train_set, batch_size)
test_loader = DataLoader(test_set, batch_size, shuffle=False)

model = MLP((hidden_size, hidden_size, 10), activation=relu)
optimizer = Adam().setup(model)
optimizer.add_hook(WeightDecay(1e-4))  # Weight decay

if cuda.gpu_enable:
    train_loader.to_gpu()
    test_loader.to_gpu()
    model.to_gpu()

# 加载参数
# if os.path.exists('mnist_weights.npz'):
#     model.load_weights('mnist_weights.npz')

for epoch in range(max_epoch):
    sum_loss, sum_acc = 0, 0

    for x, t in train_loader:
        y = model(x)
        loss = softmax_cross_entropy(y, t)
        acc = accuracy(y, t)
        model.cleargrads()
        loss.backward()
        optimizer.update()

        sum_loss += float(loss.data) * len(t)
        sum_acc += float(acc.data) * len(t)

    print('epoch: {}'.format(epoch+1))
    print('train loss: {}, accuracy: {}'.format(
        sum_loss / len(train_set), sum_acc / len(train_set)))

    sum_loss, sum_acc = 0, 0
    with no_grad():
        for x, t in test_loader:
            y = model(x)
            loss = softmax_cross_entropy(y, t)
            acc = accuracy(y, t)
            sum_loss += float(loss.data) * len(t)
            sum_acc += float(acc.data) * len(t)

    print('test loss: {}, accuracy: {}'.format(
        sum_loss / len(test_set), sum_acc / len(test_set)))

# model.save_weights('mnist_weights.npz')
